﻿Imports System.IO

Public Class main

    Private wimFile As String
    Private mountLocation As String
    Private wimMounted As Boolean = False
    Private indexList As New List(Of wimInfo)
    Private selectedIndex As New wimInfo
    Private ProcExitCode As String
    Private driversPath As String
    Private errList As New List(Of String)
    Private lastDISMcmd As String

    Private Enum AppStates
        Idle
        Busy
    End Enum

    Private _OutText As OutText
    Delegate Sub OutText(ByVal line As String)

    Private Sub main_Load(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles MyBase.Load
        SplitContainer1.Panel1MinSize = 400
        SplitContainer1.Panel2MinSize = 200
        _OutText = New OutText(AddressOf MyOutText)
        CheckMount()
        If wimMounted = True And wimFile = "" Then
            reloadMountInfo()
        End If
    End Sub

    Private Sub reloadMountInfo()
        SetState(AppStates.Busy)
        Dim reloadMount As New Process
        With reloadMount.StartInfo
            .RedirectStandardOutput = True
            .RedirectStandardError = True
            .UseShellExecute = False
            .CreateNoWindow = True
            .FileName = "dism.exe"
            .Arguments = "/Get-mountedWimInfo"
        End With
        reloadMount.Start()
        While Not reloadMount.StandardOutput.EndOfStream
            Dim line = Split(reloadMount.StandardOutput.ReadLine, " : ")

            Select Case line(0)
                Case "Mount Dir"
                    txtMountLocation.Text = line(1)
                    mountLocation = line(1)
                Case "Image File"
                    txtWIMfile.Text = line(1)
                    wimFile = line(1)
                    loadlvIndexs(wimFile)
                Case "Image Index"
                    selectedIndex = indexList.Find(Function(c) c.index = line(1))
                    For Each prop As System.Reflection.PropertyInfo In selectedIndex.GetType().GetProperties
                        With lvIndexDetail.Items.Add(prop.Name)
                            .SubItems.Add(prop.GetValue(selectedIndex, Nothing))
                        End With
                    Next
                    For Each col As ColumnHeader In lvIndexDetail.Columns
                        col.Width = -2
                    Next
                
            End Select
        End While
        SetState(AppStates.Idle)

    End Sub

    Private Sub SetState(ByVal newstate As AppStates)
        Select Case newstate
            Case AppStates.Idle
                Cursor = Cursors.Default
                If wimMounted = True Then
                    btnUnmount.Enabled = True
                    btnMount.Enabled = False
                    btn_dlgMountLocation.Enabled = False
                    btn_dlgWIMfile.Enabled = False
                    txtMountLocation.Enabled = False
                    txtWIMfile.Enabled = False
                    lvIndexs.Enabled = False
                    tabDriver.Enabled = True
                ElseIf wimMounted = False Then
                    btnUnmount.Enabled = False
                    btnMount.Enabled = True
                    btn_dlgMountLocation.Enabled = True
                    btn_dlgWIMfile.Enabled = True
                    txtMountLocation.Enabled = True
                    txtWIMfile.Enabled = True
                    lvIndexs.Enabled = True
                    tabDriver.Enabled = False
                End If

            Case AppStates.Busy
                Cursor = Cursors.WaitCursor
                btnMount.Enabled = False
                btnUnmount.Enabled = False
                btn_dlgMountLocation.Enabled = False
                btn_dlgWIMfile.Enabled = False
                txtMountLocation.Enabled = False
                txtWIMfile.Enabled = False
                lvIndexs.Enabled = False

        End Select

    End Sub

#Region "WIM Tab"
    Private Sub btn_dlgWIMfile_Click(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles btn_dlgWIMfile.Click
        dlgWIMfile.ShowDialog()
        GetwimFile(dlgWIMfile.FileName)
    End Sub

    Private Sub btn_dlgMountLocation_Click(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles btn_dlgMountLocation.Click
        dlgMountLocation.ShowDialog()
        GetMountDir(dlgMountLocation.SelectedPath)
    End Sub

    Private Function GetwimFile(ByVal filePath As String)
        Dim firstLine As String
        If Not FileIO.FileSystem.FileExists(filePath) Then firstLine = "File Not Found" Else firstLine = New StreamReader(filePath).ReadLine
        If Not firstLine.Contains("MSWIM") Then
            Dim msgResponse = MsgBox("Invalid or Corrupted WIM file!!" & vbCrLf & "Please select a valid wim.")
            wimFile = ""
            txtWIMfile.Focus()
            Return 1
        End If
        wimFile = filePath
        txtWIMfile.Text = wimFile
        loadlvIndexs(wimFile)
        Return 0
    End Function

    Private Function GetMountDir(ByVal folderPath As String)
        If Not Directory.Exists(folderPath) Then
            Dim msgResponse = MsgBox("This will create the directory:" & vbCrLf & folderPath, MsgBoxStyle.OkCancel, "Folder Does Not Exist")
            Select Case msgResponse
                Case MsgBoxResult.Ok
                    Directory.CreateDirectory(folderPath)
                Case MsgBoxResult.Cancel
                    mountLocation = Nothing
                    txtMountLocation.Focus()
                    Return 1
            End Select
        End If
        If Directory.GetFiles(folderPath).Count > 0 Or Directory.GetDirectories(folderPath).Count > 0 Then
            Dim msgResponse = MsgBox("Microsoft requiers an empty mount folder!!" & vbCrLf & "Files in the directory will be deleted!!" & vbCrLf & "Do you want to use this directory?", MsgBoxStyle.OkCancel, "!!!DIRECTORY HAS FILES!!!")
            Select Case msgResponse
                Case MsgBoxResult.Ok
                    For Each subdirectory In Directory.GetDirectories(folderPath)
                        Directory.Delete(subdirectory, True)
                    Next
                    For Each f In Directory.GetFiles(folderPath)
                        File.Delete(f)
                    Next
                Case MsgBoxResult.Cancel
                    mountLocation = Nothing
                    txtMountLocation.Focus()
                    Return 1
            End Select
        End If
        mountLocation = folderPath
        txtMountLocation.Text = mountLocation
        Return 0
    End Function

    Private Sub txtWIMfile_Validating(ByVal sender As Object, ByVal e As System.ComponentModel.CancelEventArgs) Handles txtWIMfile.Validating
        If txtWIMfile.Text <> "" Then GetwimFile(txtWIMfile.Text)
    End Sub

    Private Sub txtMountLocation_Validating(ByVal sender As Object, ByVal e As System.ComponentModel.CancelEventArgs) Handles txtMountLocation.Validating
        If txtMountLocation.Text <> "" Then GetMountDir(txtMountLocation.Text)
    End Sub

    Private Sub loadlvIndexs(ByVal wimFile As String)
        indexList.Clear()

        Dim wimIndex As New Process
        wimIndex.StartInfo.RedirectStandardOutput = True
        wimIndex.StartInfo.RedirectStandardError = True
        wimIndex.StartInfo.UseShellExecute = False
        wimIndex.StartInfo.CreateNoWindow = True
        wimIndex.StartInfo.FileName = "dism.exe"
        wimIndex.StartInfo.Arguments = "/Get-WimInfo /WimFile:""" & wimFile ' & """" & " /index:" & strIndex & " /MountDir:" & """" & strFolderName & """"
        wimIndex.Start()
        Dim indexCount = System.Text.RegularExpressions.Regex.Matches(wimIndex.StandardOutput.ReadToEnd, "Index : ").Count
        wimIndex.WaitForExit()

        For index = 1 To indexCount
            Dim info As New wimInfo
            Dim wimInfo As New Process()
            wimInfo.StartInfo.RedirectStandardOutput = True
            wimInfo.StartInfo.RedirectStandardError = True
            wimInfo.StartInfo.UseShellExecute = False
            wimInfo.StartInfo.CreateNoWindow = True
            wimInfo.StartInfo.FileName = "dism.exe"
            wimInfo.StartInfo.Arguments = "/Get-WimInfo /wimFile:""" & wimFile & """ /Index:" & index
            wimInfo.Start()
            While Not wimInfo.StandardOutput.EndOfStream
                Dim line = Split(wimInfo.StandardOutput.ReadLine, " : ")

                Select Case line(0)
                    Case "Index"
                        info.index = line(1)
                    Case "Name"
                        info.name = line(1)
                    Case "Description"
                        info.description = line(1)
                    Case "Size"
                        info.size = line(1)
                    Case "Architecture"
                        info.architecture = line(1)
                    Case "Hal"
                        info.hal = line(1)
                    Case "Version"
                        info.version = line(1)
                    Case "ServicePack Build"
                        info.spBuild = line(1)
                    Case "ServicePack Level"
                        info.spLevel = line(1)
                    Case "Edition"
                        info.edition = line(1)
                    Case "Installation"
                        info.installation = line(1)
                    Case "ProductType"
                        info.productType = line(1)
                    Case "ProductSuite"
                        info.productSuite = line(1)
                    Case "System Root"
                        info.systemRoot = line(1)
                    Case "Directories"
                        info.directories = line(1)
                    Case "Files"
                        info.files = line(1)
                    Case "Created"
                        info.created = line(1)
                    Case "Modified"
                        info.modified = line(1)

                End Select
            End While
            indexList.Add(info)
        Next
        lvIndexs.Items.Clear()
        If indexList.Count > 0 Then
            For Each wim As wimInfo In indexList
                With lvIndexs.Items.Add(wim.name)
                    .SubItems.Add(wim.description)
                End With
            Next
            For Each col As ColumnHeader In lvIndexs.Columns
                col.Width = -2
            Next
        End If

    End Sub

    Private Sub lvIndexs_ItemSelectionChanged(ByVal sender As Object, ByVal e As System.Windows.Forms.ListViewItemSelectionChangedEventArgs) Handles lvIndexs.ItemSelectionChanged
        If e.IsSelected Then
            lvIndexDetail.Items.Clear()
            selectedIndex = indexList.Find(Function(c) c.index = e.Item.Index + 1)
            For Each prop As System.Reflection.PropertyInfo In selectedIndex.GetType().GetProperties
                With lvIndexDetail.Items.Add(prop.Name)
                    .SubItems.Add(prop.GetValue(selectedIndex, Nothing))
                End With
            Next


        End If
        For Each col As ColumnHeader In lvIndexDetail.Columns
            col.Width = -2
        Next

    End Sub

    Private Sub Button1_Click(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles Button1.Click
        Debugger.Break()
    End Sub

    Private Sub btnMount_Click(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles btnMount.Click
        SetState(AppStates.Busy)
        Dim Arguments = "/Mount-WIM /WimFile:""" & wimFile & """" & " /index:" & selectedIndex.index & " /MountDir:" & """" & mountLocation & """"
        dismBackgroundWorker.RunWorkerAsync(Arguments)
    End Sub

    Private Sub btnUnmount_Click(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles btnUnmount.Click
        SetState(AppStates.Busy)
        Dim Arguments = "/unMount-WIM /MountDir:" & """" & mountLocation & """ /discard"
        dismBackgroundWorker.RunWorkerAsync(Arguments)

    End Sub

    Private Sub btnCleanupWim_Click(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles btnCleanupWim.Click
        Cursor = Cursors.WaitCursor
        SetState(AppStates.Busy)
        Dim Arguments = "/Cleanup-Wim"
        dismBackgroundWorker.RunWorkerAsync(Arguments)
    End Sub

    Private Sub CheckMount()
        SetState(AppStates.Busy)
        Dim mountedWIM As New Process
        With mountedWIM.StartInfo
            .RedirectStandardOutput = True
            .RedirectStandardError = True
            .UseShellExecute = False
            .CreateNoWindow = True
            .FileName = "dism.exe"
            .Arguments = "/Get-mountedWimInfo"
        End With
        mountedWIM.Start()
        Dim indexCount = System.Text.RegularExpressions.Regex.Matches(mountedWIM.StandardOutput.ReadToEnd, "No mounted images found.").Count
        mountedWIM.WaitForExit()
        If indexCount > 0 Then
            wimMounted = False
        End If

        If indexCount = 0 Then wimMounted = True
        SetState(AppStates.Idle)
    End Sub
#End Region

#Region "Drivers Tab"
    Private Sub rbDrivers_CheckedChanged(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles rbRemoveDrivers.CheckedChanged, rbAddDrivers.CheckedChanged
        If rbAddDrivers.Checked = True Then
            btnRemoveDriver.Enabled = False
            btnAddDrivers.Enabled = True
            lvinfFiles.MultiSelect = False
        End If
        If rbRemoveDrivers.Checked = True Then
            If driversPath <> mountLocation & "\windows\inf" Then getDriversPath(mountLocation & "\windows\inf")
            btnAddDrivers.Enabled = False
            btnRemoveDriver.Enabled = True
            lvinfFiles.MultiSelect = True

        End If


    End Sub

    Private Sub btn_dlgDriversFolder_Click(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles btn_dlgDriversFolder.Click
        txtDriverPath.Enabled = True
        dlgDriverLocation.ShowDialog()
        getDriversPath(dlgDriverLocation.SelectedPath)
    End Sub

    Private Function getDriversPath(ByVal folderPath As String)
        If Not Directory.Exists(folderPath) Then
            Dim msgResponse = MsgBox("Folder does not exist!")
            Return 1
        End If
        driversPath = folderPath
        txtDriverPath.Text = driversPath
        txtDriverPath.ReadOnly = True
        loadINF(driversPath)
        Return 0
    End Function

    Private Sub txtDriverPath_MouseDoubleClick(ByVal sender As Object, ByVal e As System.Windows.Forms.MouseEventArgs) Handles txtDriverPath.MouseDoubleClick
        txtDriverPath.ReadOnly = False
        txtDriverPath.SelectAll()
    End Sub

    Private Sub txtDriverPath_Validating(ByVal sender As Object, ByVal e As System.ComponentModel.CancelEventArgs) Handles txtDriverPath.Validating
        If txtDriverPath.Text <> "" And txtDriverPath.Enabled Then getDriversPath(txtDriverPath.Text)
    End Sub

    Private Sub loadINF(ByVal folderPath As String)
        Dim searchString = ""
        If rbAddDrivers.Checked Then searchString = "*.inf"
        If rbRemoveDrivers.Checked Then searchString = "oem*.inf"

        Dim infFiles As List(Of String) = Directory.GetFiles(folderPath, searchString, SearchOption.AllDirectories).ToList
        lvinfFiles.Items.Clear()
        For Each f In infFiles
            With lvinfFiles.Items.Add(f.Substring(InStrRev(f, "\")))
                .SubItems.Add(f.Substring(0, InStrRev(f, "\")))
            End With
        Next
        For Each col As ColumnHeader In lvinfFiles.Columns
            col.Width = -2
        Next

    End Sub

    Private Sub cbSelectAll_CheckedChanged(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles cbSelectAll.CheckedChanged
        If cbSelectAll.Checked Then
            For i = 0 To lvinfFiles.Items.Count - 1
                lvinfFiles.Items(i).Selected = True
            Next i
            lvinfFiles.Enabled = False
        End If
        If Not cbSelectAll.Checked Then
            For i = 0 To lvinfFiles.Items.Count - 1
                lvinfFiles.Items(i).Selected = False
            Next i
            lvinfFiles.Enabled = True
        End If
        lvinfFiles.Refresh()

    End Sub

    Private Sub btnAddDrivers_Click(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles btnAddDrivers.Click
        SetState(AppStates.Busy)
        Dim Arguments = "/image:""" & mountLocation & """ /add-driver"
        If Not cbSelectAll.Checked Then
            For Each item As ListViewItem In lvinfFiles.SelectedItems
                Arguments += " /driver:""" & item.SubItems(1).Text & item.Text & """"
            Next
        End If
        If cbSelectAll.Checked Then
            Arguments += " /driver:""" & driversPath & """ /recurse"
        End If

        'MsgBox(Arguments)
        dismBackgroundWorker.RunWorkerAsync(Arguments)

    End Sub

    Private Sub btnRemoveDriver_Click(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles btnRemoveDriver.Click
        SetState(AppStates.Busy)
        Dim Arguments = "/image:""" & mountLocation & """ /remove-driver"
        For Each item As ListViewItem In lvinfFiles.SelectedItems
            Arguments += " /driver:" & item.Text
        Next
        'MsgBox(Arguments)
        dismBackgroundWorker.RunWorkerAsync(Arguments)

    End Sub

#End Region


    Private Sub dismBackgroundWorker_DoWork(ByVal sender As System.Object, ByVal e As System.ComponentModel.DoWorkEventArgs) Handles dismBackgroundWorker.DoWork
        BeginInvoke(_OutText, "Running DISM " & e.Argument)
        Dim proc As New Process
        With proc.StartInfo
            .RedirectStandardError = True
            .RedirectStandardOutput = True
            .UseShellExecute = False
            .CreateNoWindow = True
            .FileName = "dism.exe"
            .Arguments = e.Argument
        End With

        proc.Start()
        While Not proc.StandardOutput.EndOfStream
            BeginInvoke(_OutText, proc.StandardOutput.ReadLine)
        End While
        ProcExitCode = proc.ExitCode
    End Sub

    Private Sub dismBackgroundWorker_ProgressChanged(ByVal sender As Object, ByVal e As System.ComponentModel.ProgressChangedEventArgs) Handles dismBackgroundWorker.ProgressChanged

    End Sub

    Private Sub dismBackgroundWorker_RunWorkerCompleted(ByVal sender As Object, ByVal e As System.ComponentModel.RunWorkerCompletedEventArgs) Handles dismBackgroundWorker.RunWorkerCompleted
        Cursor = Cursors.Default
        SetState(AppStates.Idle)
        'MyOutText(ProcExitCode)
        CheckErr()
        CheckMount()
        If InStr(lastDISMcmd, "/remove-driver") Then loadINF(driversPath)
    End Sub

    Private Sub CheckErr()
        If InStr(lastDISMcmd, "/remove-driver") Then

        End If

    End Sub

    Private Sub MyOutText(ByVal line As String)
        If InStr(line, "Running DISM ") Then
            lastDISMcmd = line.Replace("Running DISM ", "")
            errList.Clear()
        End If

        If InStr(line.ToLower, "error") Then
            errList.Add(line)
        End If

        tbCmdOutput.Text += vbCrLf & line
        tbCmdOutput.SelectionStart = tbCmdOutput.Text.Length
        tbCmdOutput.ScrollToCaret()
        Application.DoEvents()

    End Sub


End Class